import os
from argparse import ArgumentParser
import numpy as np
import torch.nn as nn
import torch
from PIL import Image
from matplotlib import pyplot as plt
import matplotlib

from Models_GAN import Generator
import Dataloader_Uv
import platform

CKPT_DIR = 'models'
G_FN = 'gan_g.pth'
MAX_SEQ_LEN = 100
FILENAME = 'sample.jpeg'
data_min =  3220.9690
data_max = 771885
SPLIT = '\\'

def generate(n):
    ''' Sample MIDI from trained generator model
    '''
    # prepare model
    num_feats = 1

    use_gpu = torch.cuda.is_available()
    g_model = Generator(num_feats, use_cuda=use_gpu)

    if not use_gpu:
        ckpt = torch.load(os.path.join(CKPT_DIR, G_FN), map_location='cpu')
    else:
        ckpt = torch.load(os.path.join(CKPT_DIR, G_FN))

    g_model.load_state_dict(ckpt)

    # generate from model then save 
    g_states = g_model.init_hidden(1)   #init_hidden()的参数为real_batch_size 这里默认为1
    z = torch.empty([1, MAX_SEQ_LEN-1, num_feats]).uniform_() # random vector
    if use_gpu:
        z = z.cuda()
        g_model.cuda()

    g_model.eval()

    pic_list_data = []
    for i in range(n):#将多个序列进行拼接
        g_feats, g_states = g_model(z, g_states)
        pic_data = g_feats.squeeze().cpu()
        pic_data = pic_data.detach().numpy() 
        pic_list_data.append(pic_data)

    if len(pic_list_data) > 1:
        pic_list_data = np.concatenate(pic_list_data, axis=0)
    else:
        pic_list_data = pic_list_data[0]
    return pic_list_data

def Anti_NormalizeDelta(data):
    data = data * max_100
    data = data + mean_100
    return data

def print_in_plot(deltadata,filename,SEQ_LEN):
    plt.clf()
    deltadata = deltadata.squeeze()
    range = deltadata.size(0)
    x = np.arange(0,range)
    y = deltadata.reshape([-1])
    plt.title("SEQ_LEN = " + str(SEQ_LEN)) 
    plt.xlabel("Frame") 
    plt.ylabel("Bright delta") 
    plt.plot(x,y)

    DIR = "Generator_Out" + SPLIT + "SEQ_LEN"+str(SEQ_LEN)
    if not os.path.exists(DIR):
        os.makedirs(DIR)
    plt.savefig(DIR + SPLIT + "plt_" + filename + ".png")

if __name__ == "__main__":

    if(platform.system()=='Windows'):
        SPLIT = '\\'
    else:
        SPLIT = '/'

    ARG_PARSER = ArgumentParser()
    # # number of times to execute generator model;
    # # all generated data are concatenated to form a single longer sequence
    ARG_PARSER.add_argument('--n', default=1, type=int)
    ARG_PARSER.add_argument('--filename', default='sample', type=str)
    ARG_PARSER.add_argument('--MAX_SEQ_LEN',default=20,type=int)
    ARGS = ARG_PARSER.parse_args()

    G_FN = 'gan_g_len' + str(ARGS.MAX_SEQ_LEN) + '.pth'
    FILENAME = ARGS.filename
    MAX_SEQ_LEN = ARGS.MAX_SEQ_LEN
    data = generate(ARGS.n)
    data = torch.from_numpy(data)
    dataloader = Dataloader_Uv.DataLoader(MAX_SEQ_LEN = ARGS.MAX_SEQ_LEN)
    data = dataloader.Anti_NormalizeDelta(data)

    print_in_plot(data,ARGS.filename,ARGS.MAX_SEQ_LEN)